﻿/* 
 * Copyright 2010 Justin Pealing
 * 
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 *     
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License. 
 */

using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.IO;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading;
using Microsoft.Win32.SafeHandles;

namespace Minnow.Web.Native
{
    /// <summary>
    /// Provides the ability to start processes.
    /// </summary>
    /// <remarks>
    /// This class works around some of the limitations of the .Net process implementation, namely environment variables
    /// being lower case (fixed in .Net v4.0) and the ability to pass arbitrary handles as StdInput, StdOutput and StdError.
    /// </remarks>
    internal class CreateProcess
    {
        private SafeFileHandle processHandle;

        /// <summary>
        /// Initialises a new instance of <see cref="CreateProcess" />
        /// </summary>
        /// <param name="ProcessHandle">Handle to the process.</param>
        /// <param name="StandardError">The redirected standard error stream, or null if the standard input was not redirected.</param>
        /// <param name="StandardInput">The redirected standard input stream, or null if the standard input was not redirected.</param>
        /// <param name="StandardOutput">The redirected standard output stream, or null if the standard input was not redirected.</param>
        /// <exception cref="ArgumentNullException">ProcessHandle is null.</exception>
        /// <exception cref="ArgumentException">ProcessHandle is invalid.</exception>
        private CreateProcess(SafeFileHandle ProcessHandle, StreamWriter StandardInput, StreamReader StandardOutput, StreamReader StandardError)
        {
            if (ProcessHandle == null)
            {
                throw new ArgumentNullException("ProcessHandle");
            }
            if (ProcessHandle.IsInvalid)
            {
                throw new ArgumentException("ProcessHandle is invalid");
            }

            this.processHandle = ProcessHandle;
            this.standardInput = StandardInput;
            this.standardOutput = StandardOutput;
            this.standardError = StandardError;
        }

        /// <summary>
        /// Starts the process with the parameters specified by the StartInfo property.
        /// </summary>
        /// <param name="StartInfo">Process start information.</param>
        public static CreateProcess Start(ProcessStartInfo StartInfo)
        {
            StreamWriter standardInput;
            StreamReader standardOutput;
            StreamReader standardError;
            NativeMethods.STARTUPINFO startupInfo = GetStartupInfo(
                StartInfo, 
                out standardInput, 
                out standardOutput, 
                out standardError);
            NativeMethods.PROCESS_INFORMATION processInformation;

            // TODO: Better notification of failures (such as the executable not existing)
            if (!NativeMethods.CreateProcess(
                null,
                StartInfo.CommandLine,
                null,
                null,
                true,
                GetCreationFlags(StartInfo),
                GetEnvironmentBlock(StartInfo),
                StartInfo.WorkingDirectory,
                startupInfo,
                out processInformation))
            {
                throw new Win32Exception(Marshal.GetLastWin32Error());
            }
            // TODO: This should probably be in a finally block or something
            startupInfo.Dispose();

            if (processInformation.hThread != IntPtr.Zero && processInformation.hThread != NativeMethods.INVALID_HANDLE_VALUE)
            {
                SafeFileHandle threadHandle = new SafeFileHandle(processInformation.hThread, true);
                threadHandle.Close();
            }
            if (processInformation.hProcess != IntPtr.Zero && processInformation.hProcess != NativeMethods.INVALID_HANDLE_VALUE)
            {
                return new CreateProcess(
                    new SafeFileHandle(processInformation.hProcess, true),
                    standardInput,
                    standardOutput,
                    standardError);
            }

            // TODO: Handle this better
            throw new Exception();
        }

        /// <summary>
        /// Gets the value that the associated process specified when it terminated.
        /// </summary>
        /// <exception cref="InvalidOperationException">The process has not exited.</exception>
        public int ExitCode
        {
            get
            {
                this.Refresh();
                if (!this.HasExited)
                {
                    throw new InvalidOperationException("The process has not yet exited.");
                }
                return this.exitCode;
            }
        }
        private int exitCode;

        /// <summary>
        /// Gets a value indicating whether the associated process has been terminated.
        /// </summary>
        public bool HasExited
        {
            get
            {
                this.Refresh();
                return this.hasExited;
            }
        }
        private bool hasExited;

        /// <summary>
        /// Returns the redirected standard input stream, or null if the standard input was not redirected.
        /// </summary>
        public StreamWriter StandardInput
        {
            get
            {
                return this.standardInput;
            }
        }
        private StreamWriter standardInput;

        /// <summary>
        /// Returns the redirected standard output stream, or null if the standard output was not redirected.
        /// </summary>
        public StreamReader StandardOutput
        {
            get
            {
                return this.standardOutput;
            }
        }
        private StreamReader standardOutput;

        /// <summary>
        /// Returns the redirected standard error stream, or null if the standard error was not redirected.
        /// </summary>
        public StreamReader StandardError
        {
            get
            {
                return this.standardError;
            }
        }
        private StreamReader standardError;

        /// <summary>
        /// Waits indefinitely for the associated process to exit.
        /// </summary>
        public void WaitForExit()
        {
            this.InternalWaitForExit(NativeMethods.INFINITE);
        }

        /// <summary>
        /// Waits the specified number of milliseconds for the associated process to exit.
        /// </summary>
        /// <param name="Milliseconds">The amount of time, in milliseconds, to wait for the associated process to exit.</param>
        public void WaitForExit(int Milliseconds)
        {
            this.InternalWaitForExit((uint)Milliseconds);
        }

        /// <summary>
        /// Immediately stops the associated process.
        /// </summary>
        public void Kill()
        {
            throw new NotImplementedException();
        }

        /// <summary>
        /// Waits the specified number of milliseconds for the associated process to exit.
        /// </summary>
        /// <param name="Milliseconds">The amount of time, in milliseconds, to wait for the associated process to exit.</param>
        private void InternalWaitForExit(uint Milliseconds)
        {
            // TODO: Handle the case where the process has already exited.
            NativeMethods.WaitForSingleObject(this.processHandle, (uint)Milliseconds);
        }

        /// <summary>
        /// Refreshes HasExited and ExitCodes properties.
        /// </summary>
        private void Refresh()
        {
            if (this.hasExited)
            {
                return;
            }

            if (NativeMethods.GetExitCodeProcess(this.processHandle, out this.exitCode))
            {
                this.hasExited = this.exitCode != NativeMethods.STILL_ACTIVE;
            }
        }

        /// <summary>
        /// Creates an anonymous pipe, suitable for use with the StartInfo standard input for a new process
        /// </summary>
        /// <param name="ParentWriter">Writer to be used by the parent process to write to the pipe.</param>
        /// <returns>Read handle to be passed to the child process.</returns>
        private static SafeFileHandle CreateInputPipe(out StreamWriter ParentWriter)
        {
            SafeFileHandle readHandle;
            SafeFileHandle writeHandle;
            CreateChildPipe(out readHandle, out writeHandle, true);

            ParentWriter = new StreamWriter(new FileStream(writeHandle, FileAccess.Write, 0x1000, false), 
                Encoding.GetEncoding(NativeMethods.GetConsoleCP()), 0x1000);
            ParentWriter.AutoFlush = true;

            return readHandle;
        }

        /// <summary>
        /// Creates a handle to the write end of an anonymous pipe, suitable for use with the StartInfo standard
        /// output and standard error.
        /// </summary>
        /// <param name="ParentReader">Writer suitable for reading the read end of the pipe.</param>
        /// <returns>Write handle to be passed to the child process.</returns>
        private static SafeFileHandle CreateOutputPipe(out StreamReader ParentReader)
        {
            SafeFileHandle writeHandle;
            SafeFileHandle readHandle;
            CreateChildPipe(out writeHandle, out readHandle, false);

            ParentReader = new StreamReader(
                new FileStream(readHandle, FileAccess.Read, 0x1000, false),
                Encoding.GetEncoding(NativeMethods.GetConsoleCP()), 
                true, 
                0x1000);

            return writeHandle;
        }

        /// <summary>
        /// Creates an anonymous pipe suitable to be passed to a child process.
        /// </summary>
        /// <param name="ChildHandle">Handle to be passed to the child.</param>
        /// <param name="ParentHandle">Handle to the parent end of the pipe.</param>
        /// <param name="ChildIsRead">If true, indicates that the child handle should be the read end of the pipe.</param>
        private static void CreateChildPipe(out SafeFileHandle ChildHandle, out SafeFileHandle ParentHandle, bool ChildIsRead)
        {
            SafeFileHandle tempParentPipe;

            NativeMethods.SECURITY_ATTRIBUTES securityAttributes = new NativeMethods.SECURITY_ATTRIBUTES();
            securityAttributes.bInheritHandle = true;

            bool returnFlag;
            if (ChildIsRead)
            {
                returnFlag = NativeMethods.CreatePipe(out ChildHandle, out tempParentPipe, securityAttributes, 0);
            }
            else
            {
                returnFlag = NativeMethods.CreatePipe(out tempParentPipe, out ChildHandle, securityAttributes, 0);
            }
            if (!returnFlag || ChildHandle.IsInvalid || tempParentPipe.IsInvalid)
            {
                throw new Win32Exception();
            }

            // TODO: Close handles returned by GetCurrentProcess
            if (!NativeMethods.DuplicateHandle(
                NativeMethods.GetCurrentProcess(),
                tempParentPipe,
                NativeMethods.GetCurrentProcess(),
                out ParentHandle,
                0,
                false,
                NativeMethods.DUPLICATE_SAME_ACCESS))
            {
                throw new Win32Exception(Marshal.GetLastWin32Error());
            }
            tempParentPipe.Close();
        }

        /// <summary>
        /// Returns a new STARTUPINFO object for use with CreateProcess.
        /// </summary>
        /// <returns>New STARTUPINFO object.</returns>
        private static NativeMethods.STARTUPINFO GetStartupInfo(ProcessStartInfo StartInfo, out StreamWriter StandardInput, out StreamReader StandardOutput, out StreamReader StandardError)
        {
            NativeMethods.STARTUPINFO returnValue = new NativeMethods.STARTUPINFO();
            StandardInput = null;
            StandardOutput = null;
            StandardError = null;

            // TODO: Validation of parameters (especially handles)
            returnValue.dwFlags = NativeMethods.STARTF_USESTDHANDLES;
            if (StartInfo.RedirectStandardInput)
            {
                returnValue.hStdInput = CreateInputPipe(out StandardInput);
            }
            else
            {
                returnValue.hStdInput = StartInfo.StandardInputHandle == null ? 
                    new SafeFileHandle(NativeMethods.INVALID_HANDLE_VALUE, false) :
                    StartInfo.StandardInputHandle;
            }

            if (StartInfo.RedirectStandardOutput)
            {
                returnValue.hStdOutput = CreateOutputPipe(out StandardOutput);
            }
            else
            {
                returnValue.hStdOutput = StartInfo.StandardOutputHandle == null ?
                    new SafeFileHandle(NativeMethods.INVALID_HANDLE_VALUE, false) :
                    StartInfo.StandardOutputHandle;
            }

            if (StartInfo.RedirectStandardError)
            {
                returnValue.hStdError = CreateOutputPipe(out StandardError);
            }
            else
            {
                returnValue.hStdError = StartInfo.StandardErrorHandle == null ?
                    new SafeFileHandle(NativeMethods.INVALID_HANDLE_VALUE, false) :
                    StartInfo.StandardErrorHandle;
            }

            return returnValue;
        }

        /// <summary>
        /// Constructs an environment block based on the environment variables of the
        /// StartInfo property.
        /// </summary>
        /// <returns>Pointer to the environment block.</returns>
        private static IntPtr GetEnvironmentBlock(ProcessStartInfo StartInfo)
        {
            IntPtr returnValue;

            if (StartInfo.EnvironmentVariables.Count == 0)
            {
                returnValue = IntPtr.Zero;
            }
            else
            {
                byte[] bytes = CreateEnvironmentBlock(StartInfo.EnvironmentVariables);
                GCHandle handle = GCHandle.Alloc(bytes, GCHandleType.Pinned);
                returnValue = handle.AddrOfPinnedObject();
            }

            return returnValue;
        }

        /// <summary>
        /// Gets the creation flags suitable for use with CreateProcess
        /// </summary>
        /// <returns></returns>
        private static uint GetCreationFlags(ProcessStartInfo StartInfo)
        {
            uint returnValue = 0;
            returnValue |= NativeMethods.CREATE_UNICODE_ENVIRONMENT;

            if (StartInfo.CreateNoWindow)
            {
                returnValue |= NativeMethods.CREATE_NO_WINDOW;
            }

            return returnValue;
        }

        /// <summary>
        /// Constructs an environment block from the list of environment variables.
        /// </summary>
        /// <param name="EnvironmentVariables"></param>
        /// <returns></returns>
        private static byte[] CreateEnvironmentBlock(IDictionary<string, string> EnvironmentVariables)
        {
            List<string> sortedNames = new List<string>(EnvironmentVariables.Keys);
            sortedNames.Sort();

            StringBuilder returnValue = new StringBuilder();
            foreach (string name in sortedNames)
            {
                returnValue.Append(name);
                returnValue.Append("=");
                returnValue.Append(EnvironmentVariables[name]);
                returnValue.Append('\0');
            }
            returnValue.Append('\0');

            // TODO: Is this always the correct encoding?
            return Encoding.Unicode.GetBytes(returnValue.ToString());
        }
    }
}
